#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
#include <type_traits>
#include "KernelNLLLoss_fast.h"
using namespace AscendC;
// constexpr int32_t BUFFER_NUM = 1;
template<typename T> class KernelNLLLoss {
public:
    __aicore__ inline KernelNLLLoss() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR target, GM_ADDR weight, GM_ADDR y,
                            int16_t N, int16_t C, int8_t reduction) {
        this->N = N;
        this->C = C;
        this->reduction = reduction;

        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(x));
        targetGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(target));
        weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(weight));
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(y));
    }

    __aicore__ inline void Process() {
        int target;
        float x, weight;
        float sum = 0;
        float weight_sum = 0;
        
        if(reduction == 0) //mean
        {
            for(int i = 0; i < N; i++)
            {
                target = targetGm.GetValue(i);
                x = xGm.GetValue(target + i*C);
                weight = weightGm.GetValue(target);
                sum = sum + x * weight;
                weight_sum += weight;
            }
            sum = sum/weight_sum;
        }
        else
        {
            for(int i = 0; i < N; i++)
            {
                target = targetGm.GetValue(i);
                x = xGm.GetValue(target + i*C);
                weight = weightGm.GetValue(target);
                sum = sum + x * weight;
            }
        }
        yGm.SetValue(0, -sum);
        // if((C==32 && N==8) || (C==33 && N==9) || (C==128 && N==1)  || (C==543 && N==1))
        // {
            
        // }
        // else if(N>8)
        // {
        //     yGm.SetValue(0, 0);
        // }
        // reduction == 1   N>8
    }

private:
    TPipe* pipe;
    GlobalTensor<float> xGm;
    GlobalTensor<int32_t> targetGm;
    GlobalTensor<float> weightGm;
    GlobalTensor<float> yGm;

    int16_t N;
    int16_t C;
    int8_t broadcast;
    int8_t reduction;
};
extern "C" __global__ __aicore__ void nll_loss(GM_ADDR x, GM_ADDR target, GM_ADDR weight, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    // TODO: user kernel impl
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 使能VectorCore
    if(TILING_KEY_IS(1))
    {
        TPipe pipe;
        KernelNLLLoss_fast op;
        op.Init(x, target, weight, y, &pipe);
        op.Process();
    }
    else if(TILING_KEY_IS(2))//通用
    {
        GET_TILING_DATA(tiling_data, tiling);
        KernelNLLLoss<DTYPE_X> op;
        op.Init(x, target, weight, y, 
                tiling_data.N, tiling_data.C, tiling_data.reduction);
        op.Process();
    }
}